#    Authors:    Chao Li, Handing Wang, Jun Zhang, Wen Yao, Tingsong Jiang
#    Xidian University, China
#    Defense Innovation Institute, Chinese Academy of Military Science, China
#    EMAIL:      lichaoedu@126.com, hdwang@xidian.edu.cn
#    DATE:       February 2022
# ------------------------------------------------------------------------
# This code is part of the program that produces the results in the following paper:
#
# Chao Li, Handing Wang, Jun Zhang, Wen Yao, Tingsong Jiang, An Approximated Gradient Sign Method Using Differential Evolution For Black-box Adversarial Attack, IEEE Transactions on Evolutionary Computation, 2022.
#
# You are free to use it for non-commercial purposes. However, we do not offer any forms of guanrantee or warranty associated with the code. We would appreciate your acknowledgement.
# ------------------------------------------------------------------------


import random
import torch
import copy
import numpy as np
from tqdm import tqdm

class FNSDE:

    def __init__(self, neter, config, isTarget=False, class_index=None):

        self.population_size = 50
        self.generations = 200
        self.F = 0.5
        self.CR = 0.6
        self.xmin = 0.0
        self.xmax = 1.0
        self.eps = config['epsilon']
        self.device = config['device']
        self.neter = neter
        self.config = config
        self.isTarget=isTarget
        self.class_index = class_index

    def init_population(self, dim):
        population = np.zeros((self.population_size, dim))
        for i in range(self.population_size):
            for j in range(dim):
                rand_value = random.random()
                population[i,j] = self.xmin + rand_value * (self.xmax-self.xmin)
        return population


    def calculate_fitness(self, taget_image, sample_adv_images, population, second_label, first_labels, dim):
        second_label = second_label
        taget_image = taget_image.cpu().detach().numpy()
        fitness = []
        function_value=np.zeros(self.population_size)

        if self.config['dataset'] == 'MNIST' or self.config['dataset'] == 'FMNIST':
            attack_direction = np.zeros((self.population_size, 1, 28, 28))
        elif self.config['dataset'] == 'ImageNet':
            attack_direction = np.zeros((self.population_size, 3, 224, 224))
        else:
            attack_direction = np.zeros((self.population_size, 3, 32, 32))

        ## attack_direction=np.zeros((self.population_size,3,32,32))
        
        for i in range(self.population_size):
            for j in range(0,dim):
                attack_direction[i,:,:,:] = attack_direction[i,:,:,:] + population[i,j] * ( sample_adv_images[j,:,:,:] - taget_image[0,:,:,:])
            attack_direction[i, :, :, :] = np.sign(attack_direction[i, :, :, :])

        for b in range(self.population_size):
            attack_image = taget_image + self.eps * attack_direction[b, :, :, :]
            attack_image = torch.from_numpy(attack_image)
            attack_image = attack_image.to(self.device)
            outputs = self.neter.net(attack_image.float())

            ## query counts
            self.neter.query += attack_image.shape[0]

            outputs = outputs.cpu().detach().numpy()

            if self.isTarget == False:
                d = outputs[0, first_labels]
                c = np.min(outputs)
                outputs.itemset(first_labels, c)
                g = np.max(outputs)
                function_value[b] = d-g
                fitness.append(function_value[b])
            else:
                g = outputs[0, self.class_index]
                c = np.min(outputs)
                outputs.itemset(self.class_index, c)
                d = np.max(outputs)
                function_value[b] = d - g
                fitness.append(function_value[b])

        return fitness

    def mutation(self, population, dim):

        Mpopulation = np.zeros((self.population_size, dim))
        for i in range(self.population_size):
            r1 = r2 = r3 = 0
            while r1 == i or r2 == i or r3 == i or r2 == r1 or r3 == r1 or r3 == r2:
                r1 = random.randint(0, self.population_size - 1)
                r2 = random.randint(0, self.population_size - 1)
                r3 = random.randint(0, self.population_size - 1)
            Mpopulation[i] = population[r1] + self.F * (population[r2] - population[r3])

            for j in range(dim):
                if self.xmin <= Mpopulation[i, j] <= self.xmax:
                    Mpopulation[i, j] = Mpopulation[i, j]
                else:
                    Mpopulation[i,j] = self.xmin + random.random() * (self.xmax - self.xmin)
        return Mpopulation

    def crossover(self, Mpopulation, population, dim):
        Cpopulation = np.zeros((self.population_size,dim))
        for i in range(self.population_size):
            for j in range(dim):
                rand_j = random.randint(0, dim - 1)
                rand_float = random.random()
                if rand_float <= self.CR or rand_j == j:
                    Cpopulation[i, j] = Mpopulation[i, j]
                else:
                    Cpopulation[i, j] = population[i, j]
        return Cpopulation

    def selection(self, taget_image, sample_adv_images, Cpopulation, population,second_label, first_labels, dim, pfitness):
        Cfitness = self.calculate_fitness(taget_image, sample_adv_images,  Cpopulation,second_label, first_labels, dim)
        for i in range(self.population_size):
            if Cfitness[i] < pfitness[i]:
                population[i] = Cpopulation[i]
                pfitness[i] = Cfitness[i]
            else:
                population[i] = population[i]
                pfitness[i] = pfitness[i] 
        return population, pfitness

    def FDE(self, taget_image, adversarial_images, second_label, first_labels):
        num = np.size(adversarial_images, 0)
        if num >= 10:
            dim = 10
            index = random.sample(range(0, num), dim)
            sample_adv_images = adversarial_images[index]
        else:
            dim = num
            sample_adv_images = adversarial_images
            # debug
            if sample_adv_images.ndim != 4:
                return taget_image

        population = self.init_population(dim)
        fitness = self.calculate_fitness(taget_image, sample_adv_images, population,second_label, first_labels, dim)
        Best_indi_index = np.argmin(fitness)
        Best_indi = population[Best_indi_index, :]
        for step in tqdm(range(self.generations)):
            if min(fitness) < 0:
                break
            Mpopulation = self.mutation(population, dim)
            Cpopulation = self.crossover(Mpopulation, population, dim)
            population, fitness = self.selection(taget_image, sample_adv_images, Cpopulation, population, second_label, first_labels, dim, fitness)
            Best_indi_index = np.argmin(fitness)
            Best_indi = population[Best_indi_index, :]

        if self.config['dataset'] == 'MNIST' or self.config['dataset'] == 'FMNIST':
            Finalattack_sign = np.zeros((1, 1, 28, 28))
        elif self.config['dataset'] == 'ImageNet':
            Finalattack_sign = np.zeros((1, 3, 224, 224))
        else:
            Finalattack_sign = np.zeros((1, 3, 32, 32))

        ## Finalattack_sign = np.zeros((1, 3, 32, 32))
        taget_image = taget_image.cpu().detach().numpy()
        for j in range(0, dim):
            Finalattack_sign[0, :, :, :] = Finalattack_sign[0, :, :, :] + Best_indi[j] * (sample_adv_images[j, :, :, :] - taget_image[0, :, :, :])
            Final_direction = np.sign(Finalattack_sign)
            final_image = taget_image + self.eps * Final_direction
            final_image = torch.from_numpy(final_image)
            final_image = final_image.float()
            final_image[0] = torch.clamp(final_image[0], self.xmin, self.xmax)
            # final_image[0, 0, :, :] = torch.clamp(final_image[0, 0, :, :], self.xmin, self.xmax)
            # final_image[0, 1, :, :] = torch.clamp(final_image[0, 1, :, :], self.xmin, self.xmax)
            # final_image[0, 2, :, :] = torch.clamp(final_image[0, 2, :, :], self.xmin, self.xmax)

        return final_image